#!/usr/bin/env python 
# -*- coding:utf-8 -*-
import binascii
from pyDiameter.pyDiaMessage import *
from scapy.all import *

def get_data(path):
    pcap_pkts = rdpcap(path)
    dict_messaage=[]
    for packet in pcap_pkts:
        packet.show()
        if packet['SCTPChunkData'] !=None:
            decode_data=packet['SCTPChunkData']
            #whether a message +
            data_begin=decode_data.beginning
            data_end=decode_data.ending
            #get data
            msgBuff=decode_data.data
            if data_begin == 1 and data_end==1:
                msg = DiaMessage()
                msg.decode(msgBuff)
                visitMessage(msg)
            # if data_begin == 1 and data_end == 0:
            #     SCTPChunkData_seq=decode_data.tsn
            #     dict_messaage[SCTPChunkData_seq] = msgBuff
            # if data_begin == 0 and data_end == 1:
            #     SCTPChunkData_seq=decode_data.tsn
            #     dict_messaage[SCTPChunkData_seq] = msgBuff

def union_msgBuff():
    dict_messaage={'id': '001', 'age': 30, 'name': 'zhangsan', 'sex': 'male'}
    dict=[]
    for key in dict_messaage:
        message_buff=dict_messaage[key]
        dict.append(message_buff)
    print(dict)


def visitAVP(avp, tab=''):
    print("=======AVP=========")
    print(tab, end='')
    print('name:  ', avp.getAVPName())
    print(tab, end='')
    print('type:  ', avp.getAVPType())
    print(tab, end='')
    print('code:  ', avp.getAVPCode())
    print(tab, end='')
    print('flags: ', avp.getAVPFlags())
    print(tab, end='')
    print('len:   ', len(avp))
    value = avp.getAVPValue()
    if avp.getAVPVSFlag():
        print(tab, end='')
        print('vendor:', avp.getAVPVendor())
    if type(value) is list:
        print(tab, end='')
        print('====>')
        for sub in value:
            visitAVP(sub, tab + '    ')
        print(tab, end='')
        print('<====')
    else:
        print(tab, end='')
        print('value: ', value)
    print(tab, end='')
    print('-------')

def visitMessage(msg):
    print("############Frame############")
    print('len:   ', len(msg))
    print('flags: ', msg.getFlags())
    print('code:  ', msg.getCommandCode())
    print('app:   ', msg.getApplicationID())
    print('hbh:   ', msg.getHBHID())
    print('e2e:   ', msg.getE2EID())
    avps = msg.getAVPs()
    for avp in avps:
        visitAVP(avp)

if __name__ == '__main__':
    # union_msgBuff()
    path= 'D:/python/diameter/data/diameter.pcap'
    get_data(path)
    # msgBuff=msgBuffs[0]+msgBuffs[1]
    # msgBuff=binascii.b2a_hex(msgBuff).decode()
    # msgBuff=bytes.fromhex(msgBuff)
    # msg = DiaMessage()
    # msg.decode(msgBuff)
    # visitMessage(msg)

# print("=========== Add an address AVP by using DiaAVPPath ===========")
# # as we known, avp1 is a group avp, only group avp could add sub-avp.
# from pyDiameter.pyDiaAVPBasicTypes import DiaAVPStr
# from pyDiameter.pyDiaAVPPath import DiaAVPPath
# from pyDiameter.pyDiaAVPTools import address_to_bytes, bytes_to_address, time_to_bytes, bytes_to_time

# newAVP = DiaAVPStr()
# newAVP.setAVPCode(257)
# newAVP.setAVPMandatoryFlag()
# newAVP.setAVPValue(address_to_bytes(('ipv6', '2000:1234:ac12:abcd:abcd:abcd:abcd:abc1')))
# avpPath = DiaAVPPath()
# avpPath.setPath('10415/1400[0]')
# msg.addAVPByPath(avpPath, newAVP)
# visitMessage(msg)

# print("=========== Remove a Feature-List AVP in Second Supported-Feature AVP ===========")
# avpPath = DiaAVPPath()
# avpPath.setPath('10415/628[1]->10415/630[0]')
# msg.removeAVPByPath(avpPath)
# visitMessage(msg)

# Try following functions and find the result by yourself.
# msg.generateHBHID()
# msg.generateE2EID()
# msg.setRequestFlag()
# msg.clearRequestFlag()
# msg.setProxyableFlag()
# msg.clearProxyableFlag()
# msg.setEBitFlag()
# msg.clearEBitFlag()
# msg.setTBitFlag()
# msg.clearTBitFlag()